# -*- coding:utf-8 -*-
# @FileName : app\services\mcp_server\mcp_client.py
# @Time     : 2025/11/28
# @Author   : 天空之城
"""手动建立连接"""
import asyncio
import os
import asyncio
from contextlib import AsyncExitStack
from typing import Optional, Any

from mcp import ClientSession
from mcp.client.sse import sse_client

from app.config import logger
# 提前导入Constant
from app.utils.Constant import Constant


class MCPClient:
    def __init__(self, server_url):
        self.server_url = server_url
        self._exit_stack = AsyncExitStack()
        self._session: Optional[ClientSession] = None
        self._is_connected = False
        self._keep_alive_task: Optional[asyncio.Task] = None
        self._connection_lock = asyncio.Lock()
        logger.info(f"【MCP客户端】初始化，服务器地址: {server_url}")

    async def connect(self):
        
        async with self._connection_lock:
            if self._is_connected:
                logger.warning(f"【MCP客户端】已经连接到: {self.server_url}")
                return

            try:
                logger.info(f"【MCP客户端】开始连接服务器: {self.server_url}")

                # 创建SSE连接
                logger.debug(f"【MCP客户端】创建SSE通道连接: {self.server_url}")
                sse_context = sse_client(self.server_url)
                self.read, self.write = await self._exit_stack.enter_async_context(sse_context)
                logger.debug(f"【MCP客户端】SSE通道连接成功: {self.server_url}")

                # 创建MCP会话
                logger.debug(f"【MCP客户端】创建MCP会话: {self.server_url}")
                self._session = await self._exit_stack.enter_async_context(ClientSession(self.read, self.write))

                # 初始化会话
                logger.debug(f"【MCP客户端】初始化MCP会话: {self.server_url}")
                await self._session.initialize()

                self._is_connected = True
                logger.info(f"【MCP客户端】连接并初始化成功: {self.server_url}")

                # 启动保活任务
                self._keep_alive_task = asyncio.create_task(self._keep_alive())

            except Exception as e:
                logger.error(f"【MCP客户端】连接失败: {self.server_url}, 错误: {str(e)}")
                await self._cleanup()
                raise

    async def disconnect(self):
        """手动断开连接"""
        async with self._connection_lock:
            if not self._is_connected:
                logger.warning(f"【MCP客户端】未连接到: {self.server_url}")
                return

            logger.info(f"【MCP客户端】手动断开连接: {self.server_url}")
            await self._cleanup()
            logger.info(f"【MCP客户端】连接已关闭: {self.server_url}")
            # 释放MCP服务器
            logger.debug(f"【MCP客户端】释放服务器资源: {self.server_url}")
            Constant.release_task_mcp_server_by_url(self.server_url)
            logger.debug(f"【MCP客户端】服务器资源已释放: {self.server_url}")

    async def _keep_alive(self):
        """保活任务，维持长连接"""
        try:
            while self._is_connected:
                # 每隔30秒发送一个心跳或检查连接状态
                await asyncio.sleep(30)
                if self._is_connected:
                    try:
                        # 可以发送一个简单的请求来保持连接活跃
                        await self.list_tools()
                        logger.debug(f"【MCP客户端】保活检查成功: {self.server_url}")
                    except Exception as e:
                        logger.warning(f"【MCP客户端】保活检查失败，尝试重连: {self.server_url}, 错误: {str(e)}")
                        await self._reconnect()
        except asyncio.CancelledError:
            logger.debug(f"【MCP客户端】保活任务被取消: {self.server_url}")
        except Exception as e:
            logger.error(f"【MCP客户端】保活任务异常: {self.server_url}, 错误: {str(e)}")

    async def _reconnect(self):
        """重连逻辑"""
        logger.info(f"【MCP客户端】尝试重新连接: {self.server_url}")
        await self._cleanup()
        await self.connect()

    async def _cleanup(self):
        """清理资源"""
        self._is_connected = False

        # 取消保活任务
        if self._keep_alive_task and not self._keep_alive_task.done():
            self._keep_alive_task.cancel()
            try:
                await self._keep_alive_task
            except asyncio.CancelledError:
                pass

        # 关闭连接栈
        await self._exit_stack.aclose()
        self._session = None
        self._keep_alive_task = None

    async def __aenter__(self):
        """兼容原有的 async with 用法"""
        await self.connect()
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        """兼容原有的 async with 用法"""
        if exc_type:
            logger.error(
                f"【MCP客户端】退出时发生异常: {self.server_url}, "
                f"异常类型: {exc_type.__name__}, 异常信息: {str(exc_val)}"
            )
        await self.disconnect()

    async def ensure_connected(self):
        """确保连接状态"""
        if not self._is_connected:
            await self.connect()

    async def list_tools(self) -> Any:
        """获取可用工具列表"""
        await self.ensure_connected()
        try:
            logger.debug(f"【MCP客户端】请求工具列表: {self.server_url}")
            result = await self._session.list_tools()
            tool_count = len(result.tools) if result and hasattr(result, 'tools') else 0
            logger.info(f"【MCP客户端】获取工具列表成功，工具数量: {tool_count}")
            return result
        except Exception as e:
            logger.error(f"【MCP客户端】获取工具列表失败: {self.server_url}, 错误: {str(e)}")
            await self._reconnect()
            raise

    async def list_resources(self) -> Any:
        """获取可用资源列表"""
        await self.ensure_connected()
        try:
            logger.debug(f"【MCP客户端】请求资源列表: {self.server_url}")
            result = await self._session.list_resources()
            resource_count = len(result.resources) if result and hasattr(result, 'resources') else 0
            logger.info(f"【MCP客户端】获取资源列表成功，资源数量: {resource_count}")
            return result
        except Exception as e:
            logger.error(f"【MCP客户端】获取资源列表失败: {self.server_url}, 错误: {str(e)}")
            await self._reconnect()
            raise

    async def list_prompts(self) -> Any:
        """获取可用提示列表"""
        await self.ensure_connected()
        try:
            logger.debug(f"【MCP客户端】请求提示列表: {self.server_url}")
            result = await self._session.list_prompts()
            prompt_count = len(result.prompts) if result and hasattr(result, 'prompts') else 0
            logger.info(f"【MCP客户端】获取提示列表成功，提示数量: {prompt_count}")
            return result
        except Exception as e:
            logger.error(f"【MCP客户端】获取提示列表失败: {self.server_url}, 错误: {str(e)}")
            await self._reconnect()
            raise

    async def call_tool(self, name: str, arguments: dict) -> Any:
        """调用指定工具"""
        await self.ensure_connected()
        try:
            logger.info(f"【MCP客户端】调用工具: {name}, 服务器: {self.server_url}")
            result = await self._session.call_tool(name, arguments)
            logger.debug(f"【MCP客户端】工具调用结果: {name}, 结果: {result}")
            return result
        except Exception as e:
            logger.error(f"【MCP客户端】工具调用失败: {name}, 服务器: {self.server_url}, 错误: {str(e)}")
            await self._reconnect()
            raise

    @property
    def is_connected(self) -> bool:
        """获取连接状态"""
        return self._is_connected

    async def close(self):
        """关闭连接（disconnect的别名）"""
        await self.disconnect()
